classdef MatDataSet<handle
    properties
        file_info = []
        selected_file = [];%string array
        data = [];
        label = [];

    end
    
    properties(Access = private)
        suffix
        trans_len
    end
    
    %% public methods
    methods(Access = public)
        % constructor
        function obj = MatDataSet(args)
            arguments
                args.selected_file = []
                args.trans_len = 0;
            end
            
            file_path = uigetdir('D:\');
            obj.suffix = 'mat';
            obj.trans_len = args.trans_len;
            obj.selected_file = args.selected_file;
            
            % 1. find all the mat files in the <file_path>
            get_all_file(obj, file_path);
            
            % 2. select file from <selected_file>
            select_file(obj)
            
            % 3. assign <data> and <label>
            get_data_and_label(obj)
                   
        end
        
        function [x, y] = get_item(obj, index)
            x = obj.data(:, :, index);
            y = obj.label(index, :);
                        
            x = translate(x, obj.trans_len);
        end
        
        function L = len(obj)
            L = size(obj.label, 1);
        end
        
        function plot(obj, index, args)
            arguments
                obj
                index
                args.pol_channel = ["hh", "vh", "hv", "vv"];
                args.scale = 'db';
            end
            % choose polarization channel to plot
            pol_index = zeros(1, 4);
            
            if ismember("hh", args.pol_channel)
                pol_index(1) = 1;
            end
            
            if ismember("vh", args.pol_channel)
                pol_index(2) = 2;
            end
            
            if ismember("hv", args.pol_channel)
                pol_index(3) = 3;
            end
            
            if ismember("vv", args.pol_channel)
                pol_index(4) = 4;
            end
            % plot
            [x, ~] = get_item(obj, index);
            x_plt = squeeze(x(:, pol_index, :));
            for j = 1: size(x_plt,3)
                switch args.scale
                    case 'db'
                        plot(db(x_plt(:, :, j)));
                    case 'abs'
                        plot(abs(x_plt(:, :, j)));
                end
                
                legend({'hh', 'vh', 'hv', 'vv'});
                setimg(gcf);
                pause(0.5)
            end
            
        end
    end
    
    %% protected methods
    methods(Access = protected)
        
        function get_data_and_label(obj)
            for i = 1:length(obj.file_info)
                filename  = fullfile(obj.file_info(i).folder, obj.file_info(i).name);
                HRRP_struct = load(filename);
                HRRP = HRRP_struct.HRRP;
                
                num_file = size(HRRP, 3);
                label_i = repmat(obj.file_info(i).encoding, num_file, 1);
                
                obj.data = cat(3, obj.data, HRRP);
                obj.label = cat(1, obj.label, label_i);
            end
        end
        
        function select_file(obj)
            
            if ~isempty(obj.selected_file)
                cnt = 1;
                while true
                    if cnt>length(obj.file_info)
                        break;
                    end
                    % check if the file_info element is in <selected_file>
                    if ~ismember(obj.file_info(cnt).name(1:end-length(obj.suffix)-1), obj.selected_file)
                        obj.file_info(cnt) = [];
                    else
                        cnt = cnt+1;
                    end
                end  
                
                if isempty(obj.file_info)
                    error('No files is selected! Please check the obj.selected_file variable.')
                end
            end
            
            
        end
        
        function  get_all_file(obj, filepath)
            obj.file_info = ...
                get_all_file_(filepath, obj.suffix);
        end
        
    end
end

%% local function
function encoding = one_hot(n)
encoding = cell(n, 1);
for i = 1:n
    zero_vec = zeros(1, n);
    zero_vec(i) = 1;
    encoding{i} = zero_vec;
end   
end

function contents = get_all_file_(filepath, suffix)
    contents = dir(filepath);
    
    len_suffix = length(char(suffix));
    % 1. delete dir 
    contents([contents.isdir]) = [];
    
    cnt = 1;
    % 2. delete useless field
    contents = rmfield(contents, {'date', 'bytes', 'datenum', 'isdir'});
    % 3. find '*.mat' file
    while true
        
        if cnt > length(contents)
            break;
        end
        
        ct_name = char(contents(cnt).name);
        
        if ~strcmp(ct_name(end-len_suffix+1: end), suffix)
            contents(cnt) = [ ];
        else
            cnt = cnt+1;
        end
        
    end
    
    if isempty(contents)
        error('No specified files in the directory!')
    else
        contents  = struct('name', {contents.name}',...
                                       'folder', {contents.folder}',...
                                       'encoding', one_hot(length(contents)));
    end
    
end

function data = translate(data, trans_len)
    batch = 3;
    if trans_len ~= 0
        for i = 1:size(data, batch)
            sample = data(:, :, i);
            data(:, :, i) = translate_(sample, trans_len);
        end
    end
end

function data = translate_(data, trans_len)
    trans = round(trans_len*2*(rand(1)-0.5));
    if trans < 0
        temp = data(end+trans+1: end, :);
        data(end+trans+1:end, :) = [];
        data = cat(1, temp, data);
    elseif trans>0
        temp = data(1:trans, :);
        data(1:trans, :) = [];
        data = cat(1, data, temp);
    end
end